#include <stdlib.h>
#include <string.h>
#include "scrypt.h"
#include "salsa20.h"

void Xor_Block(uint32_t* in_out, uint32_t* other, uint32_t length) {
    for (int k = 0; k < length; k++) {
        in_out[k] ^= other[k];
    }
}

// Convert the last 16 words (64 bytes) of the block to a little-endian number
uint32_t Integerify(uint32_t* block, uint32_t length) {
    uint32_t last_word = block[length - 16]; // Get the last 16th word
    return last_word;
}

void BlockMix(uint32_t* Y, uint32_t length) {
    uint32_t r = length / (2 * 16); // to (2*r) 16-word chunks
    uint32_t *B, *X;

    // copy Y to B
    B = (uint32_t *) malloc(length * sizeof(uint32_t));
    memcpy(B, Y, length * sizeof(uint32_t));
    // copy last chunk of B to X
    X = (uint32_t *) malloc(16 * sizeof(uint32_t));
    memcpy(X, B + length - 16, 16 * sizeof(uint32_t));

    // BlockMix
    for (int i = 0; i < 2 * r; i++) {
        Xor_Block(X, B + i * 16, 16);
        salsa20R8(X);
        if (i % 2 == 0) {
            memcpy(Y + (8*i),            X, 16 * sizeof(uint32_t));
        } else {
            memcpy(Y + (16*r + 8*(i-1)), X, 16 * sizeof(uint32_t));
        }
    }

    // free memory
    free(B);
    free(X);
}

void ROMix(uint32_t* X, uint32_t length, uint32_t iterations) {
    uint32_t** V = (uint32_t **) malloc(iterations * sizeof(uint32_t*));
    // allocate memory for V
    for (int i = 0; i < iterations; i++) {
        V[i] = (uint32_t *) malloc(length * sizeof(uint32_t));
    }
    // ROMix Phase 1
    for (int i = 0; i < iterations; i++) {
        memcpy(V[i], X, length * sizeof(uint32_t));
        BlockMix(X, length);
    }
    // ROMix Phase 2
    for (int i = 0; i < iterations; i++) {
        uint32_t j = Integerify(X, length) % iterations;
        Xor_Block(X, V[j], length);
        BlockMix(X, length);
    }
    // free memory for V
    for (int i = 0; i < iterations; i++) {
        free(V[i]);
    }
    free(V);
}

void nativeParallelROMix(uint32_t* B, uint32_t N, uint32_t r, uint32_t p) {
    uint32_t blockSize = 32 * r; //32 Words = 128 Bytes
    for (int i = 0; i < p; i++) { //TODO: ROMix could be paralleled
        uint32_t* X = B + i * blockSize;
        ROMix(X, blockSize, N);
    }
}
